# -*- coding: utf-8 -*-

import casadi
import numpy as np
import browser as b
import time
import threading

class Ellipse(object):

    def __init__(self, opt):
        self.var = ["a", "b", "yaw", "x0", "y0"]
        for v in self.var:
            setattr(self, v, opt.variable())
    
    def __call__(self, x1, y1):
        x, y = self.inv_tran(x1, y1)
        return x**2/self.a**2 + y**2/self.b**2 - 1
    
    def init(self, opt):
        opt.subject_to(self.a > 0.0)
        opt.subject_to(self.b > 0.0)
        opt.subject_to(self.a - self.b >= 0.0)
        opt.subject_to(self.yaw >= -np.pi / 2.0)
        opt.subject_to(self.yaw <= np.pi / 2.0)
        opt.set_initial(self.a, 1.0)
        opt.set_initial(self.b, 1.0)
        opt.set_initial(self.yaw, -np.pi / 2.0)
        opt.set_initial(self.x0, 0.0)
        opt.set_initial(self.y0, 0.0)
    
    def solve(self, opt):
        self.init(opt)
        opt.solver('ipopt')
        sol = opt.solve()
        self.finish(sol)
    
    def finish(self, sol):
        self.sol = sol
        for v in self.var:
            setattr(self, f"var_{v}", getattr(self, v))
            setattr(self, v, sol.value(getattr(self, f"var_{v}")))
    
    def sample(self, t=None):
        if t is None:
            t = np.arange(0.0, np.pi * 2 + 0.1, 0.1)
        x = self.a * np.cos(t)
        y = self.b * np.sin(t)
        return self.tran(x, y)
    
    def tran(self, x, y):
        cos = casadi.cos(self.yaw)
        sin = casadi.sin(self.yaw)
        xt, yt = x, y
        x = cos * xt + sin * yt
        y = -sin * xt + cos * yt
        x += self.x0
        y += self.y0
        return x, y
    
    def inv_tran(self, x1, y1):
        cos = casadi.cos(-self.yaw)
        sin = casadi.sin(-self.yaw)
        x = cos * (x1 - self.x0) + sin * (y1 - self.y0)
        y = -sin * (x1 - self.x0) + cos * (y1 - self.y0)
        return x, y
    
    def distance(self, p, e):
        xp, yp = self.inv_tran(p[0], p[1])
        xe, ye = self.inv_tran(e[0], e[1])
        xp /= self.a
        xe /= self.a
        yp /= self.b
        ye /= self.b
        A = yp - ye
        B = - (xp - xe)
        C = xp*ye - yp*xe
        d = np.sqrt(C**2/(A**2+B**2))
        return d
    
    def isCross(self, p, e):
        xp, yp = self.inv_tran(p[0], p[1])
        xe, ye = self.inv_tran(e[0], e[1])
        xp /= self.a
        xe /= self.a
        yp /= self.b
        ye /= self.b
        A = yp - ye
        B = - (xp - xe)
        C = xp*ye - yp*xe
        d = np.sqrt(C**2/(A**2+B**2))
        if d >= 1.0:
            return False
        cos = A / np.sqrt(A**2+B**2)
        sin = B / np.sqrt(A**2+B**2)
        y1 = -sin * xp + cos * yp
        y2 = -sin * xe + cos * ye
        y0 = np.sqrt(1 - d*d)
        if y1 <= -y0 and y2 <= -y0:
            return False
        if y1 >= y0 and y2 >= y0:
            return False
        return True
    
    def subjectElp(self, opt, elp):
        x0 = opt.variable()
        y0 = opt.variable()
        yaw = opt.variable()
        opt.set_initial(x0, 0.0)
        opt.set_initial(y0, 0.0)
        opt.set_initial(yaw, 0.0)
        p = [x0, y0]
        e = [x0 + np.cos(yaw), y0 + np.sin(yaw)]
        # opt.subject_to(det(elp.x0, elp.y0, p, e) * 
        #     det(self.x0, self.y0, p, e) <= 0.0)
        opt.subject_to(self.distance(p, e) >= 1.0)
        opt.subject_to(elp.distance(p, e) >= 1.0)


z = [
    [0,  0.5,  0, 0.75, 1.5, 0.5, 0],
    [0, -0.5, -2,   -3, 0.5,   1, 0]
]
z1 = np.array([])
z2 = np.array([])
for i in range(len(z[0])):
    j = (i+1) % (len(z[0])-1)
    z1 = np.hstack((z1, np.linspace(z[0][i], z[0][j], 20)))
    z2 = np.hstack((z2, np.linspace(z[1][i], z[1][j], 20)))
Z = np.array([z1 * 2.0, z2])

def ACircle(p, e, sina):
    p = np.array(p)
    e = np.array(e)
    pe = np.linalg.norm(p - e)
    r = pe * sina / (1 - sina**2)
    oe = r / sina
    if pe <= 0.0:
        pe = 0.001
    o = e + oe / pe * (p - e)
    return o, r

def inv_ACircle(e, o, sina):
    o = np.array(o)
    e = np.array(e)
    oe = np.linalg.norm(o - e)
    pe = oe * (1 - sina**2)
    if oe <= 0.0:
        oe = 0.001
    p = e + pe / oe * (o - e)
    return p

def det(x, y, P, E):
    return -(x*(P[1]-E[1])-y*(P[0]-E[0])+(P[0]*E[1]-P[1]*E[0]))

def tl(data):
    return data[0].tolist(), data[1].tolist()


g_Purser = []
g_Eid = object()
g_E = [3.0, -3.0]
g_Eid = object()

def SetPurserList(lst):
    global g_Purser
    g_Purser = lst

def GetPurserList():
    return g_Purser

def SetEscape(e):
    global g_E
    g_E = e

b.Emit("scatter", [id(g_Eid), g_E[0], g_E[1], 0xff0000])

@b.socketio.on("controlE")
def ControlE(e):
    g_E[0] += 0.1 * (e[0] - g_E[0])
    g_E[1] += 0.1 * (e[1] - g_E[1])
    b.Emit("scatter", [id(g_Eid), g_E[0], g_E[1], 0xff0000])

@b.socketio.on("getE")
def getE():
    b.Emit("lockCamera", g_E)

class Purser(object):

    def __init__(self, x, y):
        self.elp = None
        self.elpid = object()
        self.target = object()
        self.position = [x, y]
    
    def init(self, opt):
        opt.subject_to(self.elp.a <= 0.5)
        opt.subject_to(self.elp.b <= 0.35)
        opt.subject_to(self.elp.b >= 0.1)
        opt.subject_to(self.elp(*self.position) <= 0.0)
        opt.subject_to((self.elp.x0 - g_E[0])**2 + (self.elp.y0 - g_E[1])**2 >= 0.5)
        self.elp.init(opt)
        opt.set_initial(self.elp.x0, self.position[0])
        opt.set_initial(self.elp.y0, self.position[1])
    
    def getLoss(self):
        global g_Purser
        loss = self.elp.a**(-2)+self.elp.b**(-2)
        loss = loss
        loss1 = (self.elp.x0 - g_E[0])**2 + (self.elp.y0 - g_E[1])**2
        loss += loss1 / 100.0
        sina = np.sin(np.pi / len(g_Purser))
        o, r = ACircle(self.position, g_E, sina)
        do = np.array([0.0, 0.0])
        for p in g_Purser:
            if p is self:
                continue
            o1, r1 = ACircle(p.position, g_E, sina)
            oo1 = np.linalg.norm(o - o1)
            if oo1 < r1 + r:
                if oo1 <= 0.0:
                    oo1 = 0.001
                do += (o - o1) * (r1 + r) / oo1
        o += do
        p = inv_ACircle(g_E, o, sina)
        loss2 = (self.elp.x0 - p[0])**2 + (self.elp.y0 - p[1])**2
        # loss += loss2 / 300.0
        loss += loss2 / 150.0
        # loss3 = -((self.elp.x0 - self.position[0])**2 + (self.elp.y0 - self.position[1])**2)
        # loss += loss3
        return loss
    
    def subjectElp(self, elpZ, opt):
        global g_E
        print("*"*30, elpZ.isCross(self.position, g_E))
        if elpZ.isCross(self.position, g_E):
            opt.subject_to(det(elpZ.x0, elpZ.y0, self.position, g_E) * 
                det(self.elp.x0, self.elp.y0, self.position, g_E) <= 0.0)
        opt.subject_to((self.elp.x0 - self.position[0]) * (g_E[0] - self.position[0]) + 
            (self.elp.y0 - self.position[1]) * (g_E[1] - self.position[1]) >= 0.0)
        seg = 10
        if not hasattr(self, "e"):
            self.e = 0.0
        for i in range(seg):
            opt.subject_to(elpZ(*self.elp.sample(2.0 * np.pi * i / seg)) >= self.e)
        for p in g_Purser:
            opt.subject_to(self.elp(*p.position) >= 0.0)
        # self.elp.subjectElp(opt, elpZ)
    
    def newTarget(self, elpZ):
        print(self.position)
        opt = casadi.Opti()
        self.elp = Ellipse(opt)
        loss = self.getLoss()
        opt.minimize(loss)
        if type(elpZ) is not list:
            elpZ = [elpZ]
        for e in elpZ:
            self.subjectElp(e, opt)
        self.init(opt)
        opt.solver('ipopt')
        try:
            sol = opt.solve()
            self.e = 0.0
        except BaseException as e:
            print(e)
            self.elp = None
            self.e -= 0.01
            return
        self.elp.finish(sol)
        b.Emit("plot", [id(self.elpid), *tl(self.elp.sample()), 0x00ff00])
        b.Emit("scatter", [id(self.target), self.elp.x0, self.elp.y0, 0xff0000])
        b.Emit("scatter", [id(self), self.position[0], self.position[1], 0x0000ff])
    
    def move(self):
        if self.elp is None:
            return
        v = 0.5
        p = np.array(self.position)
        es = np.array([self.elp.x0, self.elp.y0])
        d = np.linalg.norm(es - p)
        if d <= 0.0:
            self.elp = None
            return
        if not hasattr(self, "t2"):
            self.t2 = time.time()
        self.t1 = self.t2
        self.t2 = time.time()
        d1 = (self.t2-self.t1) * v
        if d1 <= d:
            p = p + (es - p) / d * d1
        else:
            p = es
            self.elp = None
        self.position = p.tolist()
        b.Emit("scatter", [id(self), self.position[0], self.position[1], 0x0000ff])